package com.zoe.flume.sink.elasticsearch;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import org.apache.commons.lang.StringUtils;
import org.apache.flume.*;
import org.apache.flume.conf.Configurable;
import org.apache.flume.instrumentation.SinkCounter;
import org.apache.flume.sink.AbstractSink;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Created by shenyuhang on 2018/2/1.
 */
public class ElasticSearchSink extends AbstractSink implements Configurable {

    private static final Logger logger = LoggerFactory
            .getLogger(ElasticSearchSink.class);

    private final CounterGroup counterGroup = new CounterGroup();

    private int batchSize = ElasticSearchSinkConstants.DEFAULT_BATCH_SIZE;
    private String clusterName = ElasticSearchSinkConstants.DEFAULT_CLUSTER_NAME;
    private String[] esHostAndPorts;
    private String indexType = ElasticSearchSinkConstants.DEFAULT_INDEX_TYPE;
    private SinkCounter sinkCounter;
    private ElasticSearchTransportClient client;
    private DateTimeIndexNameBuilder indexNameBuilder;

    @Override
    public void configure(Context context) throws FlumeException {
        if (StringUtils.isNotBlank(context.getString(ElasticSearchSinkConstants.HOSTNAMES))) {
            esHostAndPorts = StringUtils.deleteWhitespace(
                    context.getString(ElasticSearchSinkConstants.HOSTNAMES)).split(",");
        }
        Preconditions.checkState(esHostAndPorts != null
                && esHostAndPorts.length > 0, "Missing Param:" + ElasticSearchSinkConstants.HOSTNAMES);



        if (StringUtils.isNotBlank(context.getString(ElasticSearchSinkConstants.INDEX_TYPE))) {
            this.indexType = context.getString(ElasticSearchSinkConstants.INDEX_TYPE);
        }

        if (StringUtils.isNotBlank(context.getString(ElasticSearchSinkConstants.CLUSTER_NAME))) {
            this.clusterName = context.getString(ElasticSearchSinkConstants.CLUSTER_NAME);
        }

        if (StringUtils.isNotBlank(context.getString(ElasticSearchSinkConstants.BATCH_SIZE))) {
            this.batchSize = Integer.parseInt(context.getString(ElasticSearchSinkConstants.BATCH_SIZE));
        }

        indexNameBuilder = new DateTimeIndexNameBuilder();
        indexNameBuilder.configure(context);

        if (sinkCounter == null) {
            sinkCounter = new SinkCounter(getName());
        }
    }

    @Override
    public Status process() throws EventDeliveryException {
        logger.debug("processing...");
        Status status = Status.READY;
        Channel channel = getChannel();
        Transaction txn = channel.getTransaction();
        try {
            txn.begin();
            int count;
            for (count = 0; count < batchSize; ++count) {
                Event event = channel.take();
                if (event == null) {
                    break;
                }
                client.setSerializer(new ElasticSearchEventSerializer()).addEvent(event);
            }

            if (count <= 0) {
                sinkCounter.incrementBatchEmptyCount();
                counterGroup.incrementAndGet("channel.underflow");
                status = Status.BACKOFF;
            } else {
                if (count < batchSize) {
                    sinkCounter.incrementBatchUnderflowCount();
                    status = Status.BACKOFF;
                } else {
                    sinkCounter.incrementBatchCompleteCount();
                }

                sinkCounter.addToEventDrainAttemptCount(count);
                client.submit();
            }
            txn.commit();
            sinkCounter.addToEventDrainSuccessCount(count);
            counterGroup.incrementAndGet("transaction.success");
        } catch (Throwable ex) {
            try {
                txn.rollback();
                counterGroup.incrementAndGet("transaction.rollback");
            } catch (Exception ex2) {
                logger.error(
                        "Exception in rollback. Rollback might not have been successful.",
                        ex2);
            }

            if (ex instanceof Error || ex instanceof RuntimeException) {
                logger.error("Failed to commit transaction. Transaction rolled back.",
                        ex);
                Throwables.propagate(ex);
            } else {
                logger.error("Failed to commit transaction. Transaction rolled back.",
                        ex);
                throw new EventDeliveryException(
                        "Failed to commit transaction. Transaction rolled back.", ex);
            }
        } finally {
            txn.close();
        }
        return status;
    }

    @Override
    public synchronized void start() {
        logger.info("ElasticSearch sink {} started");

        client = ElasticSearchTransportClient.builder()
                .setTransportClient(clusterName, esHostAndPorts)
                .setIndexName(indexNameBuilder.getIndexName())
                .setIndexType(indexType)
                .setBulk();

        sinkCounter.start();
        super.start();
    }

    @Override
    public synchronized void stop() {
        logger.info("ElasticSearch sink {} stopping");

        if (client != null) {
            client.close();
        }
        sinkCounter.incrementConnectionClosedCount();
        sinkCounter.stop();
        super.stop();
    }
}
